import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits import mplot3d 
from matplotlib import cm # 三维绘图使用的颜色

# 用来绘制各种二次曲面图
fig = plt.figure()
ax = fig.add_subplot(1,1,1, projection = "3d")

# 旋转单叶双曲面，绕z轴旋转【单叶双曲面】
# th = np.arange(0, 361, 1)
# z = np.arange(-5, 5, 0.1)
# X = np.zeros((z.size, th.size), dtype=np.float)
# Y = np.zeros_like(X)
# Z = np.zeros_like(Y)
# a, c = 4, 5
# for zi in range(z.size):
#     rxy = np.sqrt(a**2*(1+z[zi]**2/c**2))    
#     for xi in range(th.size):
#         X[zi, xi] = rxy*np.cos(th[xi]/180*np.pi)
#         Y[zi, xi] = rxy*np.sin(th[xi]/180*np.pi)
#         Z[zi, xi] = z[zi]
# ax.plot_surface(X, Y, Z)
# plt.show()

# 旋转双叶双曲面，绕x轴，只有一半【双叶双曲面】
# th = np.arange(0, 361, 1)
# y = np.arange(0, 5, 0.1)
# Y = np.zeros((th.size, y.size), dtype=np.float)
# Z = np.zeros_like(Y)
# X = np.zeros_like(Y)
# X2 = np.zeros_like(Y)
# a, c = 4, 5
# for xi in range(y.size):    
#     for yi in range(th.size):
#         Y[yi, xi] = y[xi]*np.cos(th[yi]/180*np.pi)
#         Z[yi, xi] = y[xi]*np.sin(th[yi]/180*np.pi)
#         X[yi, xi] = np.sqrt(a**2*(1+y[xi]**2/c**2))
#         # X2[yi, xi] = -np.sqrt(a**2*(1+y[xi]**2/c**2))
# ax.plot_surface(X, Y, Z)
# # ax.plot_surface(X2, Y, Z)
# plt.show()

# 椭球面，此时使用球坐标系，三维波束图可以采用此画法
# theta = np.arange(0, 361, 1)/180*np.pi
# phi = np.arange(-90, 91, 1)/180*np.pi
# X = np.zeros((theta.size, phi.size), dtype=np.float)
# Y = np.zeros_like(X)
# Z = np.zeros_like(Y)
# a, b, c = 3, 4, 5
# for i in range(phi.size):
#     for j in range(theta.size):
#         X[j, i] = a*np.cos(theta[j])*np.cos(phi[i])
#         Y[j, i] = b*np.sin(theta[j])*np.cos(phi[i])  
#         Z[j, i] = c*np.sin(phi[i])
# ax.plot_surface(X, Y, Z)
# plt.show()

# 椭圆抛物面

# 双曲抛物面（马鞍面）可以直接采用meshgrid画法


# 椭圆锥面
theta = np.arange(0, 361, 1)/180*np.pi
z = np.arange(0, 5, 0.1)
X = np.zeros((z.size, theta.size), dtype=np.float)
Y = np.zeros_like(X)
Z = np.zeros_like(X)
a, b = 4, 5
for ri in range(z.size):
    r = np.sqrt(z[ri])
    for xi in range(theta.size):
        X[ri, xi] = r*np.cos(theta[xi])*a
        Y[ri, xi] = r*np.sin(theta[xi])*b 
        Z[ri, xi] = r # 注意平方和不平方的区别
ax.plot_surface(X, Y, Z)
plt.show()

# 围绕Y轴示例，类似可围绕X轴
# x = np.arange(-5, 5, 1)
# x = np.array([-5, -4, -3, -2, -1, 0, -1, -2, -3, -4, -5]) # 可认为是个环
# y = np.arange(2, 6, 1)
# X, Y = np.meshgrid(x, y)
# # 列是Y，行是X
# Z = np.ones(X.shape, dtype=np.float)
# for j in range(y.size):
#     Z[j, :] = np.array([0, -1, -2, -2, -1, 0, 1, 2, 2, 1, 0])
# ax.plot_surface(X, Y, Z)
# plt.show()

# 列代表是Y，行代表是X
# x = 3*np.ones((10, ))
# y = np.arange(-6, 6, 1)
# X, Y = np.meshgrid(x, y)
# Z = np.zeros(X.shape, dtype=np.float)
# for j in range(x.size):
#     Z[:, j] = j
# ax.plot_surface(X, Y, Z)
# plt.show()

# plt.figure()
# plt.pcolormesh(x, y, z, shading='gouraud')
# plt.show()




